import torch #PyTorch深度学习框架，用于模型定义和推理
import numpy as np
from flask import Flask, request, jsonify #用于构建API
import os #用于文件路径处理
import logging #日志记录，用于调试和错误跟踪

# 配置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


# ----------------------------
# 1. 模型结构定义
# ----------------------------
class Phi_Net(torch.nn.Module):
    def __init__(self, options):
        super(Phi_Net, self).__init__()
        self.dim_x = options['dim_x']  # 输入特征维度，由v(3)+q(4)+pwm(4)
        self.dim_a = options['dim_a']  # 输出特征维度，对应补偿气动力3个分量

        self.fc1 = torch.nn.Linear(self.dim_x, 50)
        self.fc2 = torch.nn.Linear(50, 60)
        self.fc3 = torch.nn.Linear(60, 50)
        self.fc4 = torch.nn.Linear(50, self.dim_a - 1)  # 输出维度为 dim_a - 1 = 2
        self.activation = torch.nn.ReLU()

    #向前传播，数据流过网络的过程
    def forward(self, x):
        x = self.activation(self.fc1(x))
        x = self.activation(self.fc2(x))
        x = self.activation(self.fc3(x))
        x = self.fc4(x)
        #拼接一个全为1的维度，使最终输出维dim_a=3维
        if len(x.shape) == 1:#若输入是单样本（1维）
            return torch.cat([x, torch.ones(1, device=x.device)])#拼接1个1，输出3维
        else:#若输入是批量样本（2维，形状维[N,2]）
            return torch.cat([x, torch.ones([x.shape[0], 1], device=x.device)], dim=-1) #拼接1列1,输出[N,3]


class H_Net_CrossEntropy(torch.nn.Module):
    def __init__(self, options):
        super(H_Net_CrossEntropy, self).__init__()
        self.dim_a = options['dim_a']  # 输入维度，3维，与Phi_net输出一致
        self.num_c = options['num_c']  # 输出类别数

        self.fc1 = torch.nn.Linear(self.dim_a, 20)
        self.fc2 = torch.nn.Linear(20, self.num_c)
        self.activation = torch.nn.ReLU()

    def forward(self, a):#向前传播
        x = self.activation(self.fc1(a))
        x = self.fc2(x)
        return x


# ----------------------------
# 2. 预测器类
# ----------------------------
class NeuralFlyPredictor:
    def __init__(self, model_path, options):
        self.options = options
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") #选择设备，优先GPU,否则CPU

        #初始化模型并移动到指定设备
        self.phi_net = Phi_Net(options).to(self.device).float()
        self.h_net = H_Net_CrossEntropy(options).to(self.device).float()

        self._load_model(model_path)#加载预训练模型参数

        self.phi_net.eval()#设置为评估模式
        self.h_net.eval()

    def _load_model(self, model_path):   #加载模型参数
        try:
            #加载模型checkpoint（包含权重参数），并映射到当前设备
            checkpoint = torch.load(model_path, map_location=self.device)
            self.phi_net.load_state_dict(checkpoint['phi_net_state_dict']) #加载Phi_Net权重
            if checkpoint['h_net_state_dict'] is not None:
                self.h_net.load_state_dict(checkpoint['h_net_state_dict']) #加载H_Net权重
            self.phi_net = self.phi_net.float() #确保模型参数为float类型
            self.h_net = self.h_net.float()
            logger.info(f"模型加载成功：{model_path}")
        except Exception as e: #捕获加载错误
            logger.error(f"模型加载失败：{str(e)}")
            raise ValueError(f"模型加载失败：{str(e)}")

    def preprocess(self, input_data):
        """预处理单个样本"""
        required_keys = ['v', 'q', 'pwm'] #输入特征
        if not all(key in input_data for key in required_keys): #检查是否缺少必要的输入
            error_msg = f"输入必须包含：{required_keys}"
            logger.error(error_msg)
            raise ValueError(error_msg)

        #将输入列表转换为numpy数组
        v = np.array(input_data['v'], dtype=np.float32)
        q = np.array(input_data['q'], dtype=np.float32)
        pwm = np.array(input_data['pwm'], dtype=np.float32)

        #确保数据为2维（[样本数，特征数]）
        if v.ndim == 1:
            v = v.reshape(1, -1)
            q = q.reshape(1, -1)
            pwm = pwm.reshape(1, -1)

        #检查特征维度是否正确
        if v.shape[1] != 3:
            raise ValueError(f"v应为3维，实际为{v.shape[1]}")
        if q.shape[1] != 4:
            raise ValueError(f"q应为4维，实际为{q.shape[1]}")
        if pwm.shape[1] != 4:
            raise ValueError(f"pwm应为4维，实际为{pwm.shape[1]}")

        #拼接特征维度
        x = np.concatenate([v, q, pwm], axis=1)
        #转换为PyTorch张量，移动到指定设备
        return torch.tensor(x, dtype=torch.float32, device=self.device)

    def preprocess_batch(self, input_list):
        """批量预处理样本列表"""
        processed = []
        for data in input_list: #逐个处理样本
            processed.append(self.preprocess(data))
        return torch.cat(processed, dim=0) #拼接为一个张量（[N,11]）

    def predict_single(self, input_data, adaptinput, adaptlabel, lam=0):
        """单样本预测"""
        return self.predict_batch([input_data], adaptinput, adaptlabel, lam)[0]

    def predict_batch(self, batch_data, adaptinput, adaptlabel, lam=0):
        """批量预测多个样本"""
        try:
            # 1. 预处理适应集数据（用于计算模型适配参数a）
            adaptinput_tensor = self.preprocess_batch(adaptinput)
            adaptlabel_tensor = torch.tensor(adaptlabel, dtype=torch.float32, device=self.device)
            if adaptlabel_tensor.ndim == 1:
                adaptlabel_tensor = adaptlabel_tensor.unsqueeze(1)  # 确保形状为 (K, 3)

            # 2. 计算适应集的 Phi 和参数 a
            with torch.no_grad():
                Phi = self.phi_net(adaptinput_tensor)  # (K, dim_a)
                Phi_T = Phi.transpose(0, 1)  # (dim_a, K)
                A = torch.inverse(torch.mm(Phi_T, Phi) + lam * torch.eye(self.options['dim_a'], device=self.device))
                a = torch.mm(torch.mm(A, Phi_T), adaptlabel_tensor)  # (dim_a, 3)

            # 3. 批量预处理待预测数据并推理
            batch_tensor = self.preprocess_batch(batch_data)  # (N, dim_x)
            with torch.no_grad():
                phi_outputs = self.phi_net(batch_tensor)  # (N, dim_a)
                fa_preds = torch.matmul(phi_outputs, a)  # (N, 3)
                fa_np = fa_preds.cpu().numpy()

            # 4. 转换结果格式
            return [fa.tolist() for fa in fa_np]
        except Exception as e:
            logger.error(f"批量预测出错：{str(e)}")
            raise


# ----------------------------
# 3. API服务（支持批量预测接口）
# ----------------------------
app = Flask(__name__)


def get_default_options(): #定义模型配置参数
    return {
        'dim_x': 11,  #输入特征维度
        'dim_y': 3,   #输出目标维度
        'dim_a': 3,   #Phi_Net输出维度
        'num_c': 6,   #H_Net分类类别数
        'features': ['v', 'q', 'pwm'] #输入特征名称
    }

#配置模型路径和初始化预测器
MODEL_PATH = os.path.join(os.path.dirname(__file__), 'best_model.pth') #模型文件路径
options = get_default_options()
predictor = NeuralFlyPredictor(model_path=MODEL_PATH, options=options)


@app.route('/predict', methods=['POST'])
def predict_single_api():
    """单样本预测接口"""
    try:
        input_data = request.json
        if not input_data:
            return jsonify({'error': '请求体为空'}), 400

        #提取适应集数据和正则化参数
        adaptinput = input_data.get('adaptinput')
        adaptlabel = input_data.get('adaptlabel')
        lam = input_data.get('lam', 0)

        #检查适应集是否存在
        if adaptinput is None or adaptlabel is None:
            return jsonify({'error': '请求中必须包含 adaptinput 和 adaptlabel'}), 400

        #调用单样本预测，返回结果
        result = predictor.predict_single(input_data, adaptinput, adaptlabel, lam)
        return jsonify({'fa': result})
    except Exception as e:
        logger.error(f"单样本预测出错：{str(e)}")
        return jsonify({'error': str(e)}), 500


@app.route('/predict_batch', methods=['POST'])
def predict_batch_api():
    """批量预测接口（高效处理多个样本）"""
    try:
        input_data = request.json
        if not input_data:
            return jsonify({'error': '请求体为空'}), 400

        # 解析批量数据和适应集
        batch_data = input_data.get('batch_data') #待预测的批量样本列表
        adaptinput = input_data.get('adaptinput') #适应集输入列表
        adaptlabel = input_data.get('adaptlabel') #适应集标签列表
        lam = input_data.get('lam', 0)  #正则化参数

        # 验证数据格式
        if not isinstance(batch_data, list) or len(batch_data) == 0:
            return jsonify({'error': 'batch_data必须是非空列表'}), 400
        if adaptinput is None or adaptlabel is None:
            return jsonify({'error': '请求中必须包含 adaptinput 和 adaptlabel'}), 400
        if len(adaptinput) < 2:
            return jsonify({'error': '适应集至少需要2个样本'}), 400
        if len(adaptinput) != len(adaptlabel):
            return jsonify({'error': 'adaptinput和adaptlabel长度必须一致'}), 400

        # 批量预测
        results = predictor.predict_batch(batch_data, adaptinput, adaptlabel, lam)
        return jsonify({'fa_list': results})
    except Exception as e:
        logger.error(f"批量预测出错：{str(e)}")
        return jsonify({'error': str(e)}), 500


@app.route('/health', methods=['GET']) #健康检查接口，监控服务状态
def health_check():
    return jsonify({'status': 'healthy', 'model': 'neural_fly'})


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=False)